package cn.zhz.gen.support;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.map.MapUtil;
import com.baomidou.mybatisplus.core.toolkit.StringPool;
import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.baomidou.mybatisplus.generator.config.ConstVal;
import com.baomidou.mybatisplus.generator.config.INameConvert;
import com.baomidou.mybatisplus.generator.config.OutputFile;
import com.baomidou.mybatisplus.generator.config.builder.ConfigBuilder;
import com.baomidou.mybatisplus.generator.config.po.TableField;
import com.baomidou.mybatisplus.generator.config.po.TableInfo;
import com.baomidou.mybatisplus.generator.config.rules.DbColumnType;
import com.baomidou.mybatisplus.generator.config.rules.IColumnType;
import com.baomidou.mybatisplus.generator.config.rules.NamingStrategy;
import com.baomidou.mybatisplus.generator.engine.VelocityTemplateEngine;

import java.io.File;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * 代码生成器支持自定义[DTO\VO等]模版
 *
 * @author ZHZ
 * @since 2022-11-04
 */
public class EnhanceVelocityTemplateEngine extends VelocityTemplateEngine {

    private static final List<IColumnType> DATE_TYPE_LIST = CollUtil.newArrayList(DbColumnType.DATE, DbColumnType.LOCAL_DATE, DbColumnType.LOCAL_TIME, DbColumnType.LOCAL_TIME, DbColumnType.LOCAL_DATE_TIME);

    @Override
    protected void outputCustomFile(Map<String, String> customFile, TableInfo tableInfo, Map<String, Object> objectMap) {
        ConfigBuilder configBuilder = (ConfigBuilder) objectMap.get("config");
        INameConvert.DefaultNameConvert defaultNameConvert = new INameConvert.DefaultNameConvert(configBuilder.getStrategyConfig());
        // 重新获取按照不同策略生成的entityName
        String entityName = defaultNameConvert.entityNameConvert(tableInfo);
        String otherPath = this.getPathInfo(OutputFile.other);

        // 渲染公共数据
        renderData(entityName, tableInfo, objectMap);
        customFile.forEach((key, value) -> {

            // dto放一个路径下
            String dtoPackageKey = key.contains("Dto") ? "dto" : key;

            // 包配置
            Map<String, String> packageMap = (Map<String, String>) objectMap.get("package");
            objectMap.put(dtoPackageKey.concat("Package"), packageMap.get(ConstVal.OTHER).concat(StringPool.DOT).concat(dtoPackageKey));

            // 自定义class首字母大写
            String fileClassName = entityName.concat(NamingStrategy.capitalFirst(key));
            objectMap.put(key, fileClassName);
            // vo驼峰
            if (key.contains("vo")) {
                objectMap.put("voCamelName", NamingStrategy.underlineToCamel(fileClassName));
            }

            String filePath = otherPath + File.separator + dtoPackageKey + File.separator + fileClassName + suffixJavaOrKt();
            this.outputFile(new File(filePath), objectMap, value);

        });
    }

    /**
     * 渲染公共数据
     *
     * @param entityName
     * @param tableInfo
     * @param objectMap
     */
    private void renderData(String entityName, TableInfo tableInfo, Map<String, Object> objectMap) {


        // 处理vo,dto字段包
        objectMap.put("importPackages", getImportPackages(tableInfo, objectMap));
        // 处理主键字段
        TableField tableField = tableInfo.getFields().stream().filter(TableField::isKeyFlag).findFirst().orElse(null);

        if (tableInfo.isHavePrimaryKey() && tableField != null) {
            objectMap.put("idClassName", tableField.getColumnType().getType());
            objectMap.put("idPropertyCapitalName", NamingStrategy.capitalFirst(tableField.getPropertyName()));
        } else {
            objectMap.put("idClassName", "Long");
            objectMap.put("idPropertyCapitalName", "Id");
        }

        //自定义controller映射路径
        objectMap.put("controllerMappingHyphen", StringUtils.camelToHyphen(entityName));
        objectMap.put("controllerMappingEntityPath", NamingStrategy.underlineToCamel(entityName));
        // service驼峰
        objectMap.put("serviceCamelName", NamingStrategy.underlineToCamel(tableInfo.getServiceName()));
        // entity驼峰
        objectMap.put("entityCamelName", NamingStrategy.underlineToCamel(tableInfo.getEntityName()));

    }

    /**
     * 获取自定义要导入的包
     *
     * @param tableInfo
     * @param objectMap
     * @return
     */
    private Set<String> getImportPackages(TableInfo tableInfo, Map<String, Object> objectMap) {

        Set<String> importPackages = tableInfo.getFields().stream()
                .filter(tableField -> tableField.getColumnType() != null && tableField.getColumnType().getPkg() != null)
                .map(tableField -> tableField.getColumnType().getPkg())
                .collect(Collectors.toSet());

        // valid
        tableInfo.getFields().stream().filter(tableField -> !tableField.getMetaInfo().isNullable()).forEach(field -> {
            if (DbColumnType.STRING.equals(field.getColumnType())) {
                importPackages.add("javax.validation.constraints.NotBlank");
            } else {
                importPackages.add("javax.validation.constraints.NotNull");
            }
        });


        if (MapUtil.getBool(objectMap, "enableFormat", false)) {

            // long转string
            tableInfo.getFields().stream().filter(tableField -> DbColumnType.LONG.equals(tableField.getColumnType())).findFirst().ifPresent(field -> {
                importPackages.add("com.fasterxml.jackson.annotation.JsonFormat");
            });

            // 有时间类型就导入DateTimeFormat包
            tableInfo.getFields().stream().filter(tableField -> DATE_TYPE_LIST.contains(tableField.getColumnType())).findFirst().ifPresent(field -> {
                importPackages.add("com.fasterxml.jackson.annotation.JsonFormat");
                importPackages.add("org.springframework.format.annotation.DateTimeFormat");
                importPackages.add(field.getColumnType().getPkg());
            });

        }

        return importPackages;
    }


}
